import torch
import torchvision
from torch import nn


class Discriminator(nn.Module):

    def __init__(self, input_shape) -> None:
        super(Discriminator, self).__init__()

        channels, height, width = input_shape

        self.output_shape = (1, height // 2 ** 4, width // 2 ** 4)

        def discriminator_block(in_filters, out_filters, normalize=True):
            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]
            if normalize:
                layers.append(nn.InstanceNorm2d(out_filters))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers
        
        # self.main = nn.Sequential(*[
        #     nn.Conv2d(3, 64, 4, 2, 1),
        #     nn.LeakyReLU(0.2),
        #     nn.Conv2d(64, 128, 4, 2, 1),
        #     nn.InstanceNorm2d(128),
        #     nn.LeakyReLU(0.2),
        #     nn.Conv2d(128, 256, 4, 2, 1),
        #     nn.InstanceNorm2d(256),
        #     nn.LeakyReLU(0.2),
        #     nn.Conv2d(256, 512, 4, 1, 1),
        #     nn.InstanceNorm2d(512),
        #     nn.LeakyReLU(0.2),
        #     nn.Conv2d(512, 1, 4, 1, 1),
        # ])
        self.main = nn.Sequential(
            *discriminator_block(channels, 64, normalize=False),
            *discriminator_block(64, 128),
            *discriminator_block(128, 256),
            *discriminator_block(256, 512),
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(512, 1, 4, padding=1),
        )

    def forward(self, x):
        return self.main(x)
    

class ResidualBlock(nn.Module):
    def __init__(self, in_features) -> None:
        super(ResidualBlock, self).__init__()

        block = [
            nn.Conv2d(in_features, in_features, 3, 1, 1, padding_mode='reflect'),
            nn.InstanceNorm2d(in_features),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_features, in_features, 3, 1, 1, padding_mode='reflect'),
            nn.InstanceNorm2d(in_features),
        ]

        self.block = nn.Sequential(*block)

    def forward(self, x):
        return x + self.block(x)


class GeneratorResNet(nn.Module): 
    # input_shape: (3, 224, 224)
    def __init__(self, input_shape, num_residual_blocks) -> None:
        super(GeneratorResNet, self).__init__()
        
        channels = input_shape[0]
        out_features = 64
        model = [
            nn.ReflectionPad2d(channels),
            nn.Conv2d(channels, out_features, 7),
            nn.InstanceNorm2d(out_features),
            nn.ReLU(inplace=True),
        ]
        in_features = out_features # in: 64, out: 64

        for _ in range(2):
            out_features *= 2 # 1 in: 64, out: 128 2 in: 128 out: 256
            # out_features *= 4
            model += [
                nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
                nn.InstanceNorm2d(out_features),
                nn.ReLU(inplace=True),
            ]
            in_features = out_features # 1 in: 128, out: 128 2 in: 256 out: 256
        
        for _ in range(num_residual_blocks):
            model += [ResidualBlock(out_features)]

        for _ in range(2):
            out_features //= 2 # 1 in: 256, out: 128 2 in: 128 out: 64
            model += [
                nn.Upsample(scale_factor=2),
                nn.Conv2d(in_features, out_features, 3, stride=1, padding=1),
                nn.InstanceNorm2d(out_features),
                nn.ReLU(inplace=True),
            ]
            in_features = out_features # 1 in: 128, out: 128 2 in: 64, out: 64

        model += [nn.ReflectionPad2d(channels), nn.Conv2d(out_features, channels, 7), nn.Tanh()]

        self.model = nn.Sequential(*model)

    def forward(self, x):
        # print(x.shape, self.model(x).shape)
        return self.model(x)

def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
        if hasattr(m, "bias") and m.bias is not None:
            nn.init.constant_(m.bias.data, 0.0)
    elif classname.find("BatchNorm2d") != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0.0)